"""
协作服务模块
负责云端与边侧节点的协作业务逻辑，包括数据传输、模型部署、样本收集等功能
"""

import json
import asyncio
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
from datetime import datetime
from sqlalchemy.orm import Session

from database import Node, Model, TrainingTask, get_db
from config import config
from utils import LogManager
from exceptions import NotFoundException, ValidationException
from edge import EdgeClient, EdgeManager, EdgeAggregationRequest, EdgeModelDeployment, EdgeBatchOperationResult

# 初始化日志系统
LogManager.initialize_logging(config.log.level)


class CollaborateService:
    """云端与边侧节点协作服务"""
    
    def __init__(self):
        self.config = config
        self.logger = LogManager.get_logger(self.__class__.__name__)
        self.edge_client = EdgeClient(http_timeout=60)
        self.edge_manager = EdgeManager(self.edge_client)
    
    def _get_node_by_id(self, db: Session, node_id: str) -> Node:
        """获取节点对象"""
        node = db.query(Node).filter(Node.node_id == node_id).first()
        if not node:
            raise NotFoundException(f"节点 {node_id} 不存在")
        return node
    
    async def get_edge_node_health(self, db: Session, node_id: str) -> Dict[str, Any]:
        """获取边侧节点健康状态"""
        try:
            node = self._get_node_by_id(db, node_id)
            health_status = await self.edge_client.check_health(node)
            
            self.logger.info(f"获取边侧节点健康状态: {node_id} - {health_status.status}")
            return health_status.model_dump()
            
        except Exception as e:
            self.logger.error(f"获取边侧节点健康状态异常: {str(e)}")
            raise
    
    async def get_edge_samples_info(self, db: Session, node_id: str) -> Dict[str, Any]:
        """获取边侧节点样本信息"""
        try:
            node = self._get_node_by_id(db, node_id)
            samples_info = await self.edge_client.get_samples_info(node)
            
            self.logger.info(f"获取边侧节点样本信息成功: {node_id}")
            return {
                "success": True,
                "node_id": node_id,
                "data": samples_info.model_dump(),
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error(f"获取边侧节点样本信息异常: {str(e)}")
            raise
    
    async def request_edge_samples_aggregation(
        self, 
        db: Session, 
        node_id: str, 
        aggregation_params: Dict[str, Any]
    ) -> Dict[str, Any]:
        """请求边侧节点进行样本聚合"""
        try:
            node = self._get_node_by_id(db, node_id)
            
            # 构建聚合请求
            aggregation_request = EdgeAggregationRequest(
                difficulty_threshold=aggregation_params.get("difficulty_threshold"),
                time_start=aggregation_params.get("time_start"),
                time_end=aggregation_params.get("time_end"),
                max_samples=aggregation_params.get("max_samples", 1000)
            )
            
            aggregation_result = await self.edge_client.request_samples_aggregation(
                node, aggregation_request
            )
            
            self.logger.info(f"请求边侧节点样本聚合成功: {node_id}")
            return aggregation_result.model_dump()
            
        except Exception as e:
            self.logger.error(f"请求边侧节点样本聚合异常: {str(e)}")
            raise
    
    async def coordinated_samples_aggregation(
        self,
        db: Session,
        node_ids: List[str],
        aggregation_params: Dict[str, Any]
    ) -> Dict[str, Any]:
        """协调的样本聚合操作（多节点）"""
        try:
            result = await self.edge_manager.coordinated_sample_aggregation(
                db, node_ids, aggregation_params
            )
            
            self.logger.info(f"协调样本聚合完成，总样本数: {result.get('total_aggregated_samples', 0)}")
            return result
            
        except Exception as e:
            self.logger.error(f"协调样本聚合异常: {str(e)}")
            raise
    
    async def get_edge_upload_status(self, db: Session, node_id: str, task_id: str) -> Dict[str, Any]:
        """获取边侧节点上传任务状态"""
        try:
            node = self._get_node_by_id(db, node_id)
            upload_status = await self.edge_client.get_upload_status(node, task_id)
            
            self.logger.info(f"获取边侧节点上传状态成功: {node_id}, 任务: {task_id}")
            return {
                "success": True,
                **upload_status.model_dump()
            }
            
        except Exception as e:
            self.logger.error(f"获取边侧节点上传状态异常: {str(e)}")
            raise
    
    async def deploy_model_to_edge(
        self, 
        db: Session, 
        node_id: str, 
        model_id: str, 
        deployment_config: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """部署模型到边侧节点"""
        try:
            # 获取节点和模型信息
            node = self._get_node_by_id(db, node_id)
            model = db.query(Model).filter(Model.model_id == model_id).first()
            if not model:
                raise NotFoundException(f"模型 {model_id} 不存在")
            
            # 构建模型部署配置
            cloud_base_url = self.config.cloud_server.get("public_url", "http://localhost:8000")
            model_download_url = f"{cloud_base_url}/api/v1/models/{model_id}/download"
            
            deployment = EdgeModelDeployment(
                model_id=model_id,
                model_url=model_download_url,
                model_path=deployment_config.get("model_path") if deployment_config else None
            )
            
            # 部署模型
            deployment_result = await self.edge_client.deploy_model(node, deployment)
            
            if deployment_result.success:
                self.logger.info(f"模型部署成功: {model_id} -> {node_id}")
                
                # 更新节点的已部署模型列表
                deployed_models = json.loads(node.deployed_models or "[]")
                if model_id not in deployed_models:
                    deployed_models.append(model_id)
                    node.deployed_models = json.dumps(deployed_models)
                    db.commit()
                
                return {
                    "success": True,
                    "node_id": node_id,
                    "model_id": model_id,
                    "deployment_result": deployment_result.data,
                    "timestamp": deployment_result.timestamp
                }
            else:
                raise ValidationException(f"模型部署失败: {deployment_result.error}")
                
        except Exception as e:
            self.logger.error(f"部署模型到边侧节点异常: {str(e)}")
            raise
    
    async def trigger_edge_inference(self, db: Session, node_id: str) -> Dict[str, Any]:
        """触发边侧节点推理"""
        try:
            node = self._get_node_by_id(db, node_id)
            inference_result = await self.edge_client.trigger_inference(node)
            
            self.logger.info(f"触发边侧节点推理成功: {node_id}")
            return {
                "success": True,
                **inference_result.model_dump()
            }
            
        except Exception as e:
            self.logger.error(f"触发边侧节点推理异常: {str(e)}")
            raise
    
    async def get_edge_model_info(self, db: Session, node_id: str) -> Dict[str, Any]:
        """获取边侧节点模型信息"""
        try:
            node = self._get_node_by_id(db, node_id)
            model_info = await self.edge_client.get_model_info(node)
            
            self.logger.info(f"获取边侧节点模型信息成功: {node_id}")
            return {
                "success": True,
                **model_info.model_dump()
            }
          
        except Exception as e:
            self.logger.error(f"获取边侧节点模型信息异常: {str(e)}")
            raise
    
    async def get_edge_system_status(self, db: Session, node_id: str) -> Dict[str, Any]:
        """获取边侧节点系统状态"""
        try:
            node = self._get_node_by_id(db, node_id)
            system_status = await self.edge_client.get_system_status(node)
            
            self.logger.info(f"获取边侧节点系统状态成功: {node_id}")
            return {
                "success": True,
                **system_status.model_dump()
            }
            
        except Exception as e:
            self.logger.error(f"获取边侧节点系统状态异常: {str(e)}")
            raise
    
    async def request_edge_model_download(self, db: Session, node_id: str) -> Dict[str, Any]:
        """请求边侧节点从云端下载模型"""
        try:
            node = self._get_node_by_id(db, node_id)
            download_result = await self.edge_client.request_model_download(node)
            
            self.logger.info(f"请求边侧节点下载模型成功: {node_id}")
            return download_result.model_dump()
            
        except Exception as e:
            self.logger.error(f"请求边侧节点下载模型异常: {str(e)}")
            raise
    
    async def get_edge_available_models(self, db: Session, node_id: str) -> Dict[str, Any]:
        """获取边侧节点可用模型列表"""
        try:
            node = self._get_node_by_id(db, node_id)
            available_models = await self.edge_client.get_available_models(node)
            
            self.logger.info(f"获取边侧节点可用模型成功: {node_id}")
            return available_models.model_dump()
            
        except Exception as e:
            self.logger.error(f"获取边侧节点可用模型异常: {str(e)}")
            raise
    
    async def check_edge_cloud_model_status(self, db: Session, node_id: str) -> Dict[str, Any]:
        """检查边侧节点与云端模型状态"""
        try:
            node = self._get_node_by_id(db, node_id)
            cloud_status = await self.edge_client.check_cloud_model_status(node)
            
            self.logger.info(f"检查边侧节点云端模型状态成功: {node_id}")
            return cloud_status.model_dump()
            
        except Exception as e:
            self.logger.error(f"检查边侧节点云端模型状态异常: {str(e)}")
            raise
    
    async def batch_operation_on_edges(
        self, 
        db: Session, 
        operation: str, 
        node_ids: List[str], 
        params: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """在多个边侧节点上批量执行操作"""
        try:
            # 使用EdgeManager处理批量操作
            if operation == "health_check":
                result = await self.edge_manager.batch_health_check(db, node_ids)
            elif operation == "samples_info":
                result = await self.edge_manager.batch_samples_info(db, node_ids)
            elif operation == "trigger_inference":
                result = await self.edge_manager.batch_trigger_inference(db, node_ids)
            elif operation == "deploy_model" and params and "model_id" in params:
                result = await self.edge_manager.batch_model_deployment(
                    db, node_ids, params["model_id"], params.get("deployment_config")
                )
            else:
                raise ValidationException(f"不支持的批量操作: {operation}")
            
            self.logger.info(f"批量操作完成: {operation}")
            return result.model_dump()
            
        except Exception as e:
            self.logger.error(f"批量操作异常: {str(e)}")
            raise
    
    async def orchestrate_federated_training(
        self, 
        db: Session, 
        training_task_id: str, 
        edge_node_ids: List[str]
    ) -> Dict[str, Any]:
        """协调分布式模型学习训练任务"""
        try:
            # 验证训练任务和节点
            training_task = db.query(TrainingTask).filter(
                TrainingTask.task_id == training_task_id
            ).first()
            if not training_task:
                raise NotFoundException(f"训练任务 {training_task_id} 不存在")
            
            nodes = db.query(Node).filter(
                Node.node_id.in_(edge_node_ids),
                Node.node_type == 'edge'
            ).all()
            
            if not nodes:
                raise ValidationException("没有找到有效的边侧节点")
            
            self.logger.info(f"开始协调分布式模型学习: 任务 {training_task_id}, 节点数: {len(nodes)}")
            
            # 使用字典推导式简化结果收集
            participating_nodes = [node.node_id for node in nodes]
            
            # 阶段1: 并发收集样本信息
            self.logger.info("阶段1: 收集样本信息")
            data_collection_results = await self._execute_phase_operations(
                db, participating_nodes, "samples_info"
            )
            
            # 阶段2: 模型部署（如果训练完成）
            model_deployment_results = {}
            if training_task.status == "completed":
                self.logger.info("阶段2: 部署模型")
                model = db.query(Model).filter(
                    Model.training_task_id == training_task_id
                ).first()
                
                if model:
                    model_deployment_results = await self._execute_phase_operations(
                        db, participating_nodes, "deploy_model", {"model_id": model.model_id}
                    )
            
            # 阶段3: 触发推理
            self.logger.info("阶段3: 触发推理")
            inference_results = await self._execute_phase_operations(
                db, participating_nodes, "trigger_inference"
            )
            
            # 计算汇总统计
            summary = self._calculate_federated_summary(
                edge_node_ids, participating_nodes,
                data_collection_results,
                model_deployment_results,
                inference_results
            )
            
            result = {
                "training_task_id": training_task_id,
                "participating_nodes": participating_nodes,
                "data_collection_results": data_collection_results,
                "model_deployment_results": model_deployment_results,
                "inference_results": inference_results,
                "aggregation_summary": summary,
                "timestamp": datetime.now().isoformat()
            }
            
            self.logger.info(f"分布式模型学习协调完成: {summary}")
            return result
            
        except Exception as e:
            self.logger.error(f"协调分布式模型学习异常: {str(e)}")
            raise
    
    async def _execute_phase_operations(
        self, 
        db: Session, 
        node_ids: List[str], 
        operation: str, 
        params: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """执行阶段性操作"""
        results = {}
        for node_id in node_ids:
            try:
                if operation == "samples_info":
                    result = await self.get_edge_samples_info(db, node_id)
                elif operation == "deploy_model" and params:
                    result = await self.deploy_model_to_edge(
                        db, node_id, params["model_id"]
                    )
                elif operation == "trigger_inference":
                    result = await self.trigger_edge_inference(db, node_id)
                else:
                    raise ValidationException(f"不支持的操作: {operation}")
                
                results[node_id] = result
            except Exception as e:
                self.logger.error(f"节点 {node_id} 执行 {operation} 失败: {str(e)}")
                results[node_id] = {"success": False, "error": str(e)}
        
        return results
    
    def _calculate_federated_summary(
        self, 
        total_node_ids: List[str],
        participating_nodes: List[str],
        data_results: Dict,
        deployment_results: Dict,
        inference_results: Dict
    ) -> Dict[str, Any]:
        """计算分布式模型学习汇总统计"""
        successful_collections = len([
            r for r in data_results.values() if r.get("success", False)
        ])
        successful_deployments = len([
            r for r in deployment_results.values() if r.get("success", False)
        ])
        successful_inferences = len([
            r for r in inference_results.values() if r.get("success", False)
        ])
        
        total_operations = len(participating_nodes) * 2  # 数据收集 + 推理
        if deployment_results:
            total_operations += len(participating_nodes)  # + 模型部署
        
        successful_operations = successful_collections + successful_inferences
        if deployment_results:
            successful_operations += successful_deployments
        
        return {
            "total_requested_nodes": len(total_node_ids),
            "participating_nodes": len(participating_nodes),
            "successful_data_collections": successful_collections,
            "successful_model_deployments": successful_deployments,
            "successful_inferences": successful_inferences,
            "overall_success_rate": successful_operations / total_operations if total_operations > 0 else 0
        }


# 创建全局协作服务实例
collaborate_service = CollaborateService()
